import pandas as pd
import numpy as np
import os
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
import statsmodels.api as sm

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

EXCEL_PATH = "E:/M/MathModel/23国赛C题/ds/历年6月底到7月初.xlsx"
OUTPUT_DIR = "E:/M/MathModel/23国赛C题/ds/销售分析图表"


def load_and_prepare_data(file_path):
    df_2020 = pd.read_excel(file_path, sheet_name='2020')
    df_2021 = pd.read_excel(file_path, sheet_name='2021')
    df_2022 = pd.read_excel(file_path, sheet_name='2022')
    df_2023 = pd.read_excel(file_path, sheet_name='2023')

    # 合并数据
    df = pd.concat([df_2020, df_2021, df_2022, df_2023])

    # 转换日期格式
    df['销售日期'] = pd.to_datetime(df['销售日期'])

    # 修正分类名称
    df['分类名称'] = df['分类名称'].str.replace('食用菌类', '食用菌')

    # 添加工作日列 (周一到周五为工作日)
    df['is_workday'] = df['销售日期'].dt.weekday < 5

    # 添加销量(百千克)列
    df['销量(百千克)'] = df['销量(千克)'] / 100

    return df


def calculate_daily_data(df):
    # 计算每个品类每日的总销量和平均批发价
    daily_data = df.groupby(['销售日期', '分类名称']).agg(
        total_sales=('销量(百千克)', 'sum'),
        avg_wholesale=('批发价', 'mean'),
        is_workday=('is_workday', 'first')
    ).reset_index()

    # 数据清洗
    daily_data = daily_data[
        (daily_data['total_sales'] > 0) &
        (daily_data['avg_wholesale'] > 1)
        ]

    return daily_data


def perform_regression_analysis(daily_data, output_dir):
    results = []
    categories = ['花叶类', '花菜类', '辣椒类', '茄类', '水生根茎类', '食用菌']
    regression_dir = os.path.join(output_dir, "regression_results")
    os.makedirs(regression_dir, exist_ok=True)

    for category in categories:
        cat_data = daily_data[daily_data['分类名称'] == category].copy()

        # 分离工作日和非工作日数据
        workday_data = cat_data[cat_data['is_workday']]
        non_workday_data = cat_data[~cat_data['is_workday']]

        # 为每个数据集执行回归分析
        for data, data_type, label in [
            (workday_data, 'workday', '工作日'),
            (non_workday_data, 'non_workday', '非工作日')
        ]:
            if len(data) < 5:
                print(f"警告: {category} {label}数据不足({len(data)}条)，跳过回归分析")
                continue

            # 使用批发价作为自变量 (成本定价)
            X = data['avg_wholesale'].values.reshape(-1, 1)
            y = data['total_sales'].values

            # 执行线性回归
            model = LinearRegression()
            model.fit(X, y)

            # 获取回归参数
            slope = model.coef_[0]
            intercept = model.intercept_

            # 预测值
            y_pred = model.predict(X)

            # 计算模型评估指标
            r_squared = r2_score(y, y_pred)
            sse = sum((y - y_pred) ** 2)

            # 使用statsmodels获取更详细统计结果
            X_sm = sm.add_constant(X)
            model_sm = sm.OLS(y, X_sm).fit()
            p_value = model_sm.pvalues[1]  # 斜率的p值

            # 保存结果
            result = {
                'category': category,
                'data_type': data_type,
                'label': label,
                'slope': slope,
                'intercept': intercept,
                'r_squared': r_squared,
                'sse': sse,
                'p_value': p_value,
                'n_samples': len(y)
            }
            results.append(result)

            # 打印回归方程
            equation = f"Y = {slope:.4f} * X + {intercept:.4f}"
            print(f"{category} {label}回归方程: {equation}")
            print(f"  R² = {r_squared:.4f}, SSE = {sse:.2f}, p值 = {p_value:.4f}, 样本量 = {len(y)}")

            # 可视化回归结果
            plt.figure(figsize=(10, 6))
            plt.scatter(X, y, color='blue', alpha=0.7, label='实际数据')

            # 创建回归线
            x_min, x_max = X.min(), X.max()
            x_range = np.linspace(x_min, x_max, 100).reshape(-1, 1)
            y_range = model.predict(x_range)
            plt.plot(x_range, y_range, color='red', linewidth=2, label='回归线')

            plt.title(f'{category} {label}销售量与成本定价关系')
            plt.xlabel('批发价 (成本定价, 元/千克)')
            plt.ylabel('销售量 (百千克)')
            plt.legend()
            plt.grid(True, linestyle='--', alpha=0.3)

            # 添加回归方程和统计信息
            stats_text = (f'回归方程: Y = {slope:.4f}X + {intercept:.4f}\n'
                          f'R² = {r_squared:.4f}\n'
                          f'p值 = {p_value:.4f}\n'
                          f'样本量 = {len(y)}')
            plt.text(0.05, 0.95, stats_text, transform=plt.gca().transAxes,
                     fontsize=10, verticalalignment='top', bbox=dict(facecolor='white', alpha=0.8))

            plot_path = os.path.join(regression_dir, f"{category}_{data_type}_regression.png")
            plt.savefig(plot_path, dpi=300)
            plt.close()
            print(f"回归图保存至: {plot_path}")

    # 创建结果DataFrame
    results_df = pd.DataFrame(results)

    # 保存结果
    csv_path = os.path.join(regression_dir, 'regression_results.csv')
    results_df.to_csv(csv_path, index=False)
    print(f"回归结果保存至: {csv_path}")

    return results_df


def main():
    # 创建输出目录
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("步骤1: 加载和预处理数据...")
    try:
        df = load_and_prepare_data(EXCEL_PATH)
        print(f"加载成功，共 {len(df)} 条记录")
    except Exception as e:
        print(f"数据加载失败: {e}")
        return

    # 计算每日数据
    print("计算每日各品类的销售总量和平均批发价")
    daily_data = calculate_daily_data(df)
    print(f"处理完成，共 {len(daily_data)} 条每日数据")

    # 执行回归分析
    print("\n执行回归分析")
    results_df = perform_regression_analysis(daily_data, OUTPUT_DIR)

    # 输出最终结果
    print("\n回归分析结果汇总")
    for _, row in results_df.iterrows():
        print(f"\n{row['category']} {row['label']}回归结果:")
        print(f"  回归方程: Y = {row['slope']:.4f} * X + {row['intercept']:.4f}")
        print(f"  R²值: {row['r_squared']:.4f}")
        print(f"  SSE(残差平方和): {row['sse']:.4f}")
        print(f"  p值: {row['p_value']:.4f}")
        print(f"  样本数量: {row['n_samples']}")

    print("\n回归分析完成")


if __name__ == "__main__":
    main()